{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PA8SSKskdEyX"
      },
      "source": [
        "\n",
        "Copyright 2022 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "WQAPD6-tdHXF"
      },
      "outputs": [],
      "source": [
        "#@title License\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JXTGE2CLeZ5H"
      },
      "source": [
        "#  Instructions to use this colab\n",
        "\n",
        "This colab implements the [PRIME](https://arxiv.org/abs/2110.11346) model for learning conservative models for offline model-based optimization. PRIME trains a surrogate model (in this case a transformer model) to predict the objective value for a given accelerator using feasible and infeasible data. The objective to train PRIME is based on a supervised regression objective on the feasible data, and then PRIME applies an objective to push down the predicted value on infeasible points.\n",
        "\n",
        "In this colab, we implement the PRIME model, and show how it can be trained on dummy data. If you are interested in using the model and training procedure for PRIME or modifying it and want to dig into the code, please check out the section in the colab titled ``PRIME: Model Definition and Loss computation''. \n",
        "\n",
        "Running this notebook in google colab does not require installation of any dependencies since tensorflow and tensorflow_probability are installed by default, but you might need to install them if you run this notebook locally. The latest versions should work."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MrTtvT9Nc6aD"
      },
      "outputs": [],
      "source": [
        "from absl import app\n",
        "from absl import flags\n",
        "from absl import logging\n",
        "\n",
        "# These tensorflow installs are automatically provided by the\n",
        "# Google colab runtime. If you want to run this code locally,\n",
        "# make sure to install tensorflow and tensorflow_probability.\n",
        "import tensorflow.compat.v2 as tf\n",
        "import tensorflow_probability as tfp\n",
        "import numpy as np\n",
        "import os\n",
        "import pickle\n",
        "import csv\n",
        "from typing import Optional, Dict, List\n",
        "from copy import deepcopy\n",
        "\n",
        "gfile = tf.io.gfile.GFile\n",
        "\n",
        "# Default area constraint for the models we train\n",
        "AREA_THRESHOLD = 27.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cfZA-5GDE4wn"
      },
      "source": [
        "# PRIME: Model Definition and Loss computation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z6W-wFO7FK0O"
      },
      "source": [
        "## Transformer Utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2vu05EjgE8QW"
      },
      "outputs": [],
      "source": [
        "#@title Basic utility functions for training transformers\n",
        "\"\"\"\n",
        "Code largely taken from https://www.tensorflow.org/text/tutorials/transformer\n",
        "\"\"\"\n",
        "\n",
        "def get_angles(pos, i, d_model):\n",
        "  \"\"\"Get angles for using tansformer.\"\"\"\n",
        "  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))\n",
        "  return pos * angle_rates\n",
        "\n",
        "\n",
        "def positional_encoding(position, d_model):\n",
        "  \"\"\"Obtain positional encdoing for training the PRIME Transformer.\"\"\"\n",
        "  angle_rads = get_angles(np.arange(position)[:, np.newaxis],\n",
        "                          np.arange(d_model)[np.newaxis, :],\n",
        "                          d_model)\n",
        "  \n",
        "  # apply sin to even indices in the array; 2i\n",
        "  angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
        "\n",
        "  # apply cos to odd indices in the array; 2i+1\n",
        "  angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
        "\n",
        "  pos_encoding = angle_rads[np.newaxis, ...]\n",
        "\n",
        "  return tf.cast(pos_encoding, dtype=tf.float32)\n",
        "\n",
        "\n",
        "class SplitEmbeddingLayer(tf.keras.layers.Layer):\n",
        "  \"\"\"Layer for embedding individual components in a split way\"\"\"\n",
        "  def __init__(self, softmax_splits=None, output_size=32):\n",
        "    \"\"\"\n",
        "    Initialize the layer to split the input and generate embeddings for\n",
        "    each field.\n",
        "    \"\"\"\n",
        "    super(SplitEmbeddingLayer, self).__init__(trainable=True)\n",
        "    self.softmax_splits = softmax_splits\n",
        "    self.output_size = output_size\n",
        "\n",
        "    # create layers\n",
        "    self.dense_layers = []\n",
        "    print (self.softmax_splits)\n",
        "    for idx, val in enumerate(self.softmax_splits):\n",
        "      self.dense_layers.append(\n",
        "          tf.keras.layers.Dense(\n",
        "              self.output_size, name='insidelayer_' + str(idx)))\n",
        "\n",
        "    # Add position embeddings\n",
        "    self.pos_encoding = positional_encoding(position=200, d_model=output_size)\n",
        "\n",
        "  def call(self, x):\n",
        "    \"\"\"Call the Split embedding function.\"\"\"\n",
        "    split_x = tf.split(x, num_or_size_splits=self.softmax_splits, axis=-1)\n",
        "    modified_splits = []\n",
        "    idx = 0\n",
        "    for param in split_x:\n",
        "      out = self.dense_layers[int(idx)](param)\n",
        "      modified_splits.append(tf.expand_dims(out, axis=1))\n",
        "      idx += 1\n",
        "    out = tf.concat(modified_splits, axis=1)\n",
        "    # print ('Out shape before: ', out)\n",
        "    out = out + self.pos_encoding[:, :len(modified_splits), :]\n",
        "    # print ('Out shape after: ', out)\n",
        "    return out\n",
        "\n",
        "\n",
        "def scaled_dot_product_attention(q, k, v, mask):\n",
        "  \"\"\"Scaled dot product attention in transformer.\"\"\"\n",
        "  matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)\n",
        "\n",
        "  # scale matmul_qk\n",
        "  dk = tf.cast(tf.shape(k)[-1], tf.float32)\n",
        "  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)\n",
        "\n",
        "  # add the mask to the scaled tensor.\n",
        "  if mask is not None:\n",
        "    scaled_attention_logits += (mask * -1e9)\n",
        "\n",
        "  # softmax is normalized on the last axis (seq_len_k) so that the scores\n",
        "  # add up to 1.\n",
        "  attention_weights = tf.nn.softmax(\n",
        "      scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)\n",
        "  output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)\n",
        "  return output, attention_weights\n",
        "\n",
        "\n",
        "def point_wise_feed_forward_network(d_model, dff):\n",
        "  return tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(dff, activation='relu'),\n",
        "      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)\n",
        "  ])\n",
        "\n",
        "\n",
        "class MultiHeadAttention(tf.keras.layers.Layer):\n",
        "  \"\"\"Multi Head Attention for the model.\"\"\"\n",
        "\n",
        "  def __init__(self, d_model, num_heads):\n",
        "    \"\"\"Initialize the multi-head attention model.\"\"\"\n",
        "    super(MultiHeadAttention, self).__init__()\n",
        "    self.num_heads = num_heads\n",
        "    self.d_model = d_model\n",
        "\n",
        "    assert d_model % self.num_heads == 0\n",
        "\n",
        "    self.depth = d_model // self.num_heads\n",
        "\n",
        "    self.wq = tf.keras.layers.Dense(d_model)\n",
        "    self.wk = tf.keras.layers.Dense(d_model)\n",
        "    self.wv = tf.keras.layers.Dense(d_model)\n",
        "\n",
        "    self.dense = tf.keras.layers.Dense(d_model)\n",
        "\n",
        "  def split_heads(self, x, batch_size):\n",
        "    \"\"\"Split the last dimension into (num_heads, depth).\n",
        "    Transpose the result such that the\n",
        "    shape is (batch_size, num_heads, seq_len, depth)\n",
        "    \"\"\"\n",
        "    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n",
        "    return tf.transpose(x, perm=[0, 2, 1, 3])\n",
        "\n",
        "  def call(self, v, k, q, mask):\n",
        "    batch_size = tf.shape(q)[0]\n",
        "\n",
        "    q = self.wq(q)\n",
        "    k = self.wk(k)\n",
        "    v = self.wv(v)\n",
        "\n",
        "    q = self.split_heads(q, batch_size)  \n",
        "    k = self.split_heads(k, batch_size) \n",
        "    v = self.split_heads(v, batch_size)\n",
        "\n",
        "    scaled_attention, attention_weights = scaled_dot_product_attention(\n",
        "        q, k, v, mask)\n",
        "\n",
        "    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3]) \n",
        "\n",
        "    concat_attention = tf.reshape(scaled_attention,\n",
        "                                  (batch_size, -1, self.d_model)) \n",
        "\n",
        "    output = self.dense(concat_attention) \n",
        "\n",
        "    return output, attention_weights\n",
        "\n",
        "\n",
        "class TransformerLayer(tf.keras.layers.Layer):\n",
        "  \"\"\"Define the transformer layer to be used in the PRIME Transformer model.\"\"\"\n",
        "\n",
        "  def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
        "    \"\"\"Initialize the transformer layer.\"\"\"\n",
        "    super(TransformerLayer, self).__init__()\n",
        "\n",
        "    self.mha = MultiHeadAttention(d_model, num_heads)\n",
        "    self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
        "\n",
        "    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "\n",
        "    self.dropout1 = tf.keras.layers.Dropout(rate)\n",
        "    self.dropout2 = tf.keras.layers.Dropout(rate)\n",
        "\n",
        "  def call(self, x, training=True, mask=None):\n",
        "    attn_output, _ = self.mha(x, x, x, mask)  \n",
        "    # (batch_size, input_seq_len, d_model)\n",
        "    attn_output = self.dropout1(attn_output, training=training)\n",
        "    out1 = self.layernorm1(x + attn_output)  \n",
        "    # (batch_size, input_seq_len, d_model)\n",
        "\n",
        "    ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)\n",
        "    ffn_output = self.dropout2(ffn_output, training=training)\n",
        "    out2 = self.layernorm2(out1 + ffn_output) \n",
        "    # (batch_size, input_seq_len, d_model)\n",
        "    return out2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ll-Ud813QwND"
      },
      "source": [
        "## Utility and Helper Functions for Loss Computation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PCzXIHNoQtoM"
      },
      "outputs": [],
      "source": [
        "#@title Helper functions for MSE/Huber Loss computation\n",
        "\n",
        "def weighted_mse_loss(input, target, weight):\n",
        "  \"\"\"Compute weighted MSE Loss\"\"\"\n",
        "  mse_loss_val = (tf.squeeze(input) - tf.squeeze(target))**2\n",
        "  return tf.reduce_mean(mse_loss_val * tf.squeeze(weight))\n",
        "\n",
        "\n",
        "def weighted_huber_loss(input, target, weight):\n",
        "  \"\"\"Compute weighted Huber Loss\"\"\"\n",
        "  mse_loss = tf.keras.losses.Huber(\n",
        "      reduction=tf.keras.losses.Reduction.NONE)\n",
        "  return tf.reduce_mean(mse_loss(\n",
        "      y_pred=tf.squeeze(input),\n",
        "      y_true=tf.squeeze(target)) * tf.squeeze(weight))\n",
        "\n",
        "\n",
        "def weighted_approx_loss(input, target, weight):\n",
        "  \"\"\"Compute weighted Approximation Loss\"\"\"\n",
        "  abs_diff = tf.abs(tf.squeeze(input) - tf.squeeze(target))\n",
        "  ratio_diff = abs_diff / (tf.abs(tf.squeeze(target)) + 1e-6)\n",
        "  return tf.reduce_mean(ratio_diff * tf.squeeze(weight))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "b5aLcKscR0T-"
      },
      "outputs": [],
      "source": [
        "#@title Helper functions for ranking loss computation\n",
        "\n",
        "def ranking_loss(input, target, context=None):\n",
        "  \"\"\"Compute measures of ranking for the PRIMETransformerModel.\"\"\"\n",
        "  if context is not None:\n",
        "    # Compute ranking loss per context, and then average it.\n",
        "    unique_contexts, indices = tf.unique(\n",
        "        tf.squeeze(tf.cast(context, tf.int32)), name='None')\n",
        "    all_corr = []\n",
        "    for idx in range(unique_contexts.shape[0]):\n",
        "      curr_context = unique_contexts[idx]\n",
        "      locations_idx = tf.squeeze(tf.where(tf.equal(indices, curr_context)))\n",
        "      input_tmp = tf.gather(\n",
        "          tf.squeeze(input), indices=locations_idx)\n",
        "      target_tmp = tf.gather(\n",
        "          tf.squeeze(target), indices=locations_idx)\n",
        "      input_ranks = tf.argsort(input_tmp, axis=-1)\n",
        "      target_ranks = tf.argsort(target_tmp, axis=-1)\n",
        "      input_ranks = tf.cast(tf.argsort(input_ranks, axis=-1), dtype=tf.float32)\n",
        "      target_ranks = tf.cast(tf.argsort(target_ranks, axis=-1),\n",
        "                             dtype=tf.float32)\n",
        "      std_input = tf.math.reduce_std(input_ranks)\n",
        "      std_target = tf.math.reduce_std(target_ranks)\n",
        "      cov = tf.reduce_mean((target_ranks - tf.reduce_mean(target_ranks)) *\\\n",
        "                           (input_ranks - tf.reduce_mean(input_ranks)))\n",
        "      pearson_corr = cov/ (std_target * std_input)\n",
        "      all_corr.append(pearson_corr)\n",
        "    print (all_corr)\n",
        "    pearson_corr = tf.reduce_mean(pearson_corr)\n",
        "  else:\n",
        "    input = tf.squeeze(input)\n",
        "    target = tf.squeeze(target)\n",
        "    input_ranks = tf.argsort(input, axis=-1)\n",
        "    target_ranks = tf.argsort(target, axis=-1)\n",
        "    input_ranks = tf.cast(tf.argsort(input_ranks, axis=-1), dtype=tf.float32)\n",
        "    target_ranks = tf.cast(tf.argsort(target_ranks, axis=-1), dtype=tf.float32)\n",
        "    std_input = tf.math.reduce_std(input_ranks)\n",
        "    std_target = tf.math.reduce_std(target_ranks)\n",
        "    cov = tf.reduce_mean((target_ranks - tf.reduce_mean(target_ranks)) *\\\n",
        "                         (input_ranks - tf.reduce_mean(input_ranks)))\n",
        "    pearson_corr = cov/ (std_target * std_input)\n",
        "  return pearson_corr\n",
        "\n",
        "\n",
        "def ranking_trainable_loss(input, target, context=None):\n",
        "  \"\"\"Compute a differentiable ranking loss, that can be used for training.\"\"\"\n",
        "  if context is not None:\n",
        "    unique_contexts, indices = tf.unique(\n",
        "        tf.squeeze(tf.cast(context, tf.int32)), name='None')\n",
        "    all_corr = []\n",
        "    for idx in range(unique_contexts.shape[0]):\n",
        "      curr_context = unique_contexts[idx]\n",
        "      locations_idx = tf.squeeze(tf.where(tf.equal(indices, curr_context)))\n",
        "      input_tmp = tf.expand_dims(tf.gather(\n",
        "          tf.squeeze(input), indices=locations_idx), 1)\n",
        "      target_tmp = tf.expand_dims(tf.gather(\n",
        "          tf.squeeze(target), indices=locations_idx), 1)\n",
        "      input_transpose = tf.transpose(input_tmp, [1, 0]) # 1 x B\n",
        "      target_transpose = tf.transpose(target_tmp, [1, 0]) # 1 x B\n",
        "      diff_true = input_tmp - input_transpose # B x 1 - 1 x B = B x B = y_i - y_j\n",
        "      diff_pred = target_tmp - target_transpose # fx_i - fx_j\n",
        "      product = tf.sign(diff_true) * diff_pred  # sign(y_i = y_j) * (fx_i - fxj)\n",
        "      bce_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(\n",
        "          labels=tf.ones_like(product), logits=product))\n",
        "      all_corr.append(bce_loss)\n",
        "    bce_loss = tf.reduce_mean(all_corr)\n",
        "  else:\n",
        "    input_transpose = tf.transpose(input, [1, 0]) # 1 x B\n",
        "    target_transpose = tf.transpose(target, [1, 0]) # 1 x B\n",
        "    diff_true = input - input_transpose # B x 1 - 1 x B = B x B = y_i - y_j\n",
        "    diff_pred = target - target_transpose # fx_i - fx_j\n",
        "    product = tf.sign(diff_true) * diff_pred  # sign(y_i = y_j) * (fx_i - fxj)\n",
        "    bce_loss = tf.nn.sigmoid_cross_entropy_with_logits(\n",
        "        labels=tf.ones_like(product), logits=product)\n",
        "  return tf.reduce_mean(bce_loss)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2GVLIu49SUo2"
      },
      "outputs": [],
      "source": [
        "#@title Helper function for Kendall correlation\n",
        "\n",
        "def kendall_correlation(input, target, context=None):\n",
        "  \"\"\"Compute Kendall's correlation over the input, target and context.\"\"\"\n",
        "  if context is not None:\n",
        "    unique_contexts, indices = tf.unique(\n",
        "        tf.squeeze(tf.cast(context, tf.int32)), name='None')\n",
        "    all_corr = []\n",
        "    for idx in range(unique_contexts.shape[0]):\n",
        "      curr_context = unique_contexts[idx]\n",
        "      locations_idx = tf.squeeze(tf.where(tf.equal(indices, curr_context)))\n",
        "      input_tmp = tf.expand_dims(tf.gather(\n",
        "          tf.squeeze(input), indices=locations_idx), 1)\n",
        "      target_tmp = tf.expand_dims(tf.gather(\n",
        "          tf.squeeze(target), indices=locations_idx), 1)\n",
        "      input_transpose = tf.transpose(input_tmp, [1, 0])\n",
        "      target_transpose = tf.transpose(target_tmp, [1, 0])\n",
        "      diff_true = input_tmp - input_transpose\n",
        "      diff_pred = target_tmp - target_transpose\n",
        "      product = tf.sign(diff_true) * tf.sign(diff_pred)\n",
        "      positive_pairs = tf.where(tf.greater_equal(product, tf.zeros_like(product)),\n",
        "                                tf.ones_like(product), tf.zeros_like(product))\n",
        "      n = tf.cast(tf.shape(input_tmp)[0], dtype=tf.float32)\n",
        "      total_positive = tf.reduce_sum(positive_pairs) - n\n",
        "      ratio = total_positive/ (n * (n-1))\n",
        "      all_corr.append(ratio)\n",
        "    ratio = tf.reduce_mean(all_corr)\n",
        "  else:\n",
        "    input_transpose = tf.transpose(input, [1, 0])\n",
        "    target_transpose = tf.transpose(target, [1, 0])\n",
        "    diff_true = input - input_transpose\n",
        "    diff_pred = target - target_transpose\n",
        "    product = tf.sign(diff_true) * tf.sign(diff_pred)\n",
        "    positive_pairs = tf.where(tf.greater_equal(product, tf.zeros_like(product)),\n",
        "                              tf.ones_like(product), tf.zeros_like(product))\n",
        "    n = tf.cast(tf.shape(input)[0], dtype=tf.float32)\n",
        "    total_positive = tf.reduce_sum(positive_pairs) - n\n",
        "    ratio = total_positive/ (n * (n-1))\n",
        "  return 2 * ratio - 1.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CyYBvUsEGPln"
      },
      "source": [
        "## Code for the PRIME surrogate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ziT-hyEYFthV"
      },
      "outputs": [],
      "source": [
        "#@title Definition of the PRIME surrogate model, training procedure\n",
        "\n",
        "class PRIMETransformerModel(tf.keras.Model):\n",
        "  \"\"\"\n",
        "  The transformer model used by PRIME. This class implements ability to \n",
        "  instantiate a transformer model, and train it via the PRIME training objective\n",
        "  (Equation 3 in https://arxiv.org/abs/2110.11346). \n",
        "  \n",
        "  Additionally it also implements the ability to train a contextual model,\n",
        "  conditioned on the context. \n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               num_outputs,\n",
        "               num_inputs,\n",
        "               optimizer,\n",
        "               layers=(256, 256, 256),\n",
        "               penalty_weight=10.0,\n",
        "               negative_sampler=None,\n",
        "               contextual=False,\n",
        "               params_dict=None):\n",
        "    \"\"\"Initializes the PRIMETransformer model.\n",
        "\n",
        "    Args:\n",
        "      num_outputs: the dimensionality of the output of the PRIME surrogate. \n",
        "        Typically set to 1, but you can increase it to model multiple cost\n",
        "        functions together.\n",
        "      num_inputs: the dimensionality of the total number of inputs to the model.\n",
        "      optimizer: the optimizer to optimize the trainable model.\n",
        "      layers: hidden layer sizes for the feed-forward layers after extracting\n",
        "        the transformer embedding.\n",
        "      penalty_weight: the value of alpha in Equation 2 in PRIME.\n",
        "      negative_sampler: an instance of a negative sampler. A negative sampler\n",
        "        is basically an optimizer that can take in the current snapshot of the\n",
        "        this PRIMETransformerModel, and optimize the predictions of the current\n",
        "        model snapshot w.r.t its input. In the paper, we utilize an evolutionary\n",
        "        optimizer to optimize the predictions. For this code release, we present\n",
        "        a simple gradient-descent based optimizer for optimization as a\n",
        "        demonstration. Users are encouraged to pass in their relevant\n",
        "        negative sampler here.\n",
        "      contextual: bool, indicates whether we are training a contextual model\n",
        "        or a non-contextual model. Contextual is used for multi-model and\n",
        "        zero-shot experiments. \n",
        "      params_dict: dictionary. Can store additional parameters and their values.\n",
        "        This dictionary provides an easy and convenient way to add new hyper-\n",
        "        parameters, via keys of this dictionary. \n",
        "    \"\"\"\n",
        "    super().__init__()\n",
        "    self.num_inputs = num_inputs\n",
        "    self.num_outputs = num_outputs\n",
        "    self.optimizer = optimizer\n",
        "    self.params_dict = params_dict\n",
        "    self.penalty_weight = penalty_weight\n",
        "    self.contextual = contextual\n",
        "\n",
        "    # Setting the following variable to True shouldn't cause issues since\n",
        "    # it is not passed into the GradientTape, but better to be safe, and set\n",
        "    # it to false if the variable is not used.\n",
        "\n",
        "    # This variable determines the alpha multiplier in Equation 2.\n",
        "    self.log_cql_alpha = tf.Variable(tf.math.log(self.penalty_weight + 1e-6),\n",
        "                                     trainable=False)\n",
        "    self.cql_alpha_value = tf.Variable(self.penalty_weight, trainable=False)\n",
        "\n",
        "    self.negative_sampler = negative_sampler\n",
        "\n",
        "    # In the paper, we use an evolutionary optimizer for obtaining adversarial\n",
        "    # examples. However, unfortunately, this optimizer is proprietary, and so\n",
        "    # we provide the example negative sampler that uses gradient ascent, similar\n",
        "    # to conservative objective mocels https://arxiv.org/abs/2107.06882.\n",
        "    self.num_gradient_infer_steps = 0\n",
        "    if 'num_gradient_steps' in params_dict:\n",
        "      self.num_gradient_infer_steps = params_dict['num_gradient_steps']\n",
        "\n",
        "    self.opt_lr = 1e-3\n",
        "    if 'opt_lr' in params_dict:\n",
        "      self.opt_lr = params_dict['opt_lr']\n",
        "\n",
        "    # the multiplier beta in Equation 3 in the paper.\n",
        "    self.infeasible_alpha = 0.01\n",
        "    if 'infeasible_alpha' in params_dict:\n",
        "      self.infeasible_alpha = params_dict['infeasible_alpha']\n",
        "\n",
        "    # Since the input to the model is a concatenation of one-hot values\n",
        "    # representing each field, using the input_splits parameter, we partition\n",
        "    # this big input vector into a list of one-hot vectors, one corresponding \n",
        "    # to each discrete parameter. \n",
        "    self.input_splits = None\n",
        "    if 'input_splits' in params_dict:\n",
        "      self.input_splits = params_dict['input_splits']\n",
        "\n",
        "    # We use an architecture which resembles a mixture of experts, and so the\n",
        "    # following parameter decides how many parameters we wish to have.\n",
        "    self.num_votes = 1\n",
        "    if 'num_votes' in params_dict:\n",
        "      self.num_votes = params_dict['num_votes']\n",
        "\n",
        "    # Whether to add dropout or not, in intermediate layers of the model, as \n",
        "    # a means to prevent overfitting.\n",
        "    use_dropout = False\n",
        "    if 'use_dropout' in params_dict:\n",
        "      use_dropout = params_dict['use_dropout']\n",
        "\n",
        "    if self.contextual:\n",
        "      \"\"\"For contextual version of PRIME\"\"\"\n",
        "      self.num_contexts = 0\n",
        "      if 'num_contexts' in params_dict:\n",
        "        self.num_contexts = params_dict['num_contexts']\n",
        "\n",
        "    print ('Infeasible alpha: ', self.infeasible_alpha)\n",
        "    print ('CQL Alpha: ', self.log_cql_alpha)\n",
        "    print ('Num votes: ', self.num_votes)\n",
        "\n",
        "    self.input_layer = tf.keras.Input(num_inputs)\n",
        "    temp_num_inputs = num_inputs\n",
        "\n",
        "    # The following layer splits the input into a list of embeddings for \n",
        "    # each parameter. Check the SplitEmbeddingLayer class for details.\n",
        "    x = SplitEmbeddingLayer(softmax_splits=self.input_splits,\n",
        "                            output_size=64)(self.input_layer)\n",
        "    if use_dropout:\n",
        "      x = tf.keras.layers.Dropout(rate=0.1)(x)\n",
        "  \n",
        "    # Now feed the split embedding layer output into TransformerLayer\n",
        "    x = TransformerLayer(d_model=64, num_heads=8, dff=256)(x)\n",
        "    x = TransformerLayer(d_model=64, num_heads=8, dff=256)(x)\n",
        "    \n",
        "    x = tf.keras.layers.Reshape(target_shape=(640,))(x)\n",
        "    \n",
        "    if self.contextual:\n",
        "      context_input = tf.keras.Input(self.num_contexts)\n",
        "      out_context = tf.keras.layers.Dense(640, use_bias=False)(context_input)\n",
        "\n",
        "      # Pointwise multiply the contexts to make sure that the context\n",
        "      # conditioning is done properly. From https://arxiv.org/abs/1912.13465.\n",
        "      x = x * out_context\n",
        "      self._base_network = tf.keras.Model(\n",
        "        inputs=[self.input_layer, context_input], outputs=x)\n",
        "    else:\n",
        "      self._base_network = tf.keras.Model(\n",
        "          inputs=self.input_layer, outputs=x)\n",
        "\n",
        "    self.optimize_networks = [self._base_network,]\n",
        "\n",
        "    # Now feedforward layers to finish the model\n",
        "    layers = list(layers)\n",
        "    layers[0] = 64 * len(self.input_splits)\n",
        "\n",
        "    \"\"\"Voting based routing\"\"\"\n",
        "    num_networks = self.num_votes\n",
        "    self._all_networks = []\n",
        "    for jdx in range(num_networks):\n",
        "      # Make each of the networks used in routing\n",
        "      new_network = tf.keras.Sequential()\n",
        "      for idx in range(len(layers) - 1):\n",
        "        new_network.add(\n",
        "            tf.keras.layers.Dense(layers[idx+1], input_shape=(layers[idx],)))\n",
        "        new_network.add(tf.keras.layers.LeakyReLU(0.1))\n",
        "        if use_dropout:\n",
        "          new_network.add(tf.keras.layers.Dropout(rate=0.1))\n",
        "\n",
        "      new_network.add(tf.keras.layers.Dense(\n",
        "          num_outputs, input_shape=(layers[idx],)))\n",
        "      self._all_networks.append(new_network)\n",
        "\n",
        "    self.optimize_networks.extend(self._all_networks)\n",
        "\n",
        "    # Now make the network that decides the contribution of these\n",
        "    self.voting_network = tf.keras.Sequential()\n",
        "    if self.contextual:\n",
        "      self.voting_network.add(\n",
        "          tf.keras.layers.Dense(layers[1], input_shape=(2*layers[0],)))\n",
        "    else:\n",
        "      self.voting_network.add(\n",
        "          tf.keras.layers.Dense(layers[1], input_shape=(layers[0],)))\n",
        "    self.voting_network.add(tf.keras.layers.LeakyReLU(0.1))\n",
        "    if use_dropout:\n",
        "      self.voting_network.add(tf.keras.layers.Dropout(rate=0.1))\n",
        "\n",
        "    self.voting_network.add(\n",
        "        tf.keras.layers.Dense(self.num_votes, input_shape=(layers[1],)))\n",
        "\n",
        "    if self.contextual:\n",
        "      # Add the vote generation network input again\n",
        "      self.embedding_network = tf.keras.Sequential()\n",
        "      self.embedding_network.add(\n",
        "          tf.keras.layers.Dense(256))\n",
        "      self.embedding_network.add(tf.keras.layers.LeakyReLU(0.1))\n",
        "      self.embedding_network.add(\n",
        "          tf.keras.layers.Dense(layers[0]))\n",
        "\n",
        "      self.optimize_networks.append(self.embedding_network)\n",
        "    self.optimize_networks.append(self.voting_network)\n",
        "\n",
        "    print ('All networks: ', len(self.optimize_networks))\n",
        "\n",
        "  @tf.function\n",
        "  def call(self, inputs, training=True, with_logging=False):\n",
        "    \"\"\"Function to call one forward pass on the PRIME Transformer.\"\"\"\n",
        "    extra_dict = dict()\n",
        "    if not self.contextual:\n",
        "      transformer_embedding = self._base_network(inputs, training=training)\n",
        "    else:\n",
        "      # TODO(aviralkumar): Fix the hardcoded 77 input dimensionality in code\n",
        "      if not isinstance(inputs, list) and not isinstance(inputs, tuple):\n",
        "        inputs = (inputs[:, :77], inputs[:, 77:])\n",
        "\n",
        "      transformer_embedding = self._base_network(inputs, training=training)\n",
        "      \n",
        "    # Get all outputs from each expert\n",
        "    all_outputs = []\n",
        "    for idx in range(self.num_votes):\n",
        "      all_outputs.append(\n",
        "          self._all_networks[idx](transformer_embedding, training=training))\n",
        "    \n",
        "    # Get the voting probabilities\n",
        "    if self.contextual:\n",
        "      vote_input = self.embedding_network(inputs[1])\n",
        "      vote_input = tf.concat([transformer_embedding, vote_input], axis=-1)\n",
        "      vote_logit = self.voting_network(vote_input, training=training)\n",
        "    else:\n",
        "      vote_logit = self.voting_network(transformer_embedding,\n",
        "                                        training=training)\n",
        "\n",
        "    # Append all_outputs in a list and compute average score\n",
        "    all_outputs = tf.concat(all_outputs, axis=-1)   # [B x num_votes]\n",
        "    vote_prob = tf.nn.softmax(vote_logit, axis=-1)  # [B x num_votes]\n",
        "    vote_entropy = tf.reduce_sum(\n",
        "        tf.nn.log_softmax(vote_logit, axis=-1) * vote_prob, axis=-1)\n",
        "    extra_dict['vote_entropy'] = tf.reduce_mean(vote_entropy)\n",
        "    fwd_model_pred = tf.reduce_sum(vote_prob * all_outputs, axis=-1)\n",
        "    fwd_model_pred = tf.expand_dims(fwd_model_pred, axis=-1)\n",
        "    \n",
        "    if with_logging:\n",
        "      return fwd_model_pred, extra_dict\n",
        "\n",
        "    return fwd_model_pred\n",
        "\n",
        "  def compute_loss(self, data_batch, loss_type='mse', training=True,\n",
        "                   ranking_penalty_weight=0.0, inp_batch_type=None):\n",
        "    \"\"\"\n",
        "    Compute the loss function and additional logging metrics for training.\n",
        "\n",
        "    Args:\n",
        "      data_batch: A dictionary of various input fields, and their corresponding\n",
        "        tensor values. The keys for this dictionary are:\n",
        "        - design --\u003e denotes the input (accelerator config in this case)\n",
        "        - objective --\u003e denotes the objective value for the given input\n",
        "        - context_id --\u003e denotes the context vector for the case of contextual\n",
        "\n",
        "      loss_type: string, either mse or mse+rank. It essentially computes the\n",
        "        training loss used to train the PRIME model. We can optionally add some\n",
        "        ranking regularization for training if needed. Though, we did not find\n",
        "        this to be essential. \n",
        "\n",
        "      inp_batch_type: string, either 'valid' or 'mixed'. Mixed indicates that\n",
        "        the batch consists of both valid and invalid samples, whereas valid\n",
        "        indicates the samples are only valid samples.\n",
        "\n",
        "      ranking_penalty_weight: float, the weight on the ranking loss function\n",
        "        in addition to the PRIME objectives. This is not needed for PRIME, but\n",
        "        can help in some cases. So, leaving the facility here.\n",
        "    \"\"\"\n",
        "    loss_dict = dict()\n",
        "    if loss_type == 'mse':\n",
        "      fwd_loss = weighted_mse_loss\n",
        "    elif loss_type == 'mse+rank':\n",
        "      fwd_loss = weighted_mse_loss\n",
        "      ranking_loss_fn = ranking_trainable_loss\n",
        "\n",
        "    loss_dict['y_values_max'] = tf.reduce_max(data_batch['objective'])\n",
        "    loss_dict['y_values_mean'] = tf.reduce_mean(data_batch['objective'])\n",
        "\n",
        "    data_batch = data_batch.copy()\n",
        "    weights = tf.ones_like(data_batch['objective'])\n",
        "    \n",
        "    if self.contextual:\n",
        "      model_pred, extra_dict = self(\n",
        "            inputs=[data_batch['design'], data_batch['context_id']],\n",
        "            training=training, with_logging=True)\n",
        "    else:\n",
        "      model_pred, extra_dict = self(\n",
        "          data_batch['design'], training=training, with_logging=True)\n",
        "\n",
        "    loss_dict.update(extra_dict)\n",
        "\n",
        "    if self.negative_sampler is not None:\n",
        "      # This branch of the code will not run off-the-shelf, since it assumes \n",
        "      # access to a negative_sampler. A negative sampler is simply any kind of\n",
        "      # optimizer that can take in the current PRIMETransformerModel and\n",
        "      # optimize its predictions.\n",
        "      negatives_batch = self.negative_sampler.run_inference(\n",
        "          num_iters=2, model=self)\n",
        "      negatives_pred = self(inputs=negatives_batch, training=training)\n",
        "    else:\n",
        "      negatives_batch = self.infer_negatives(data_batch)\n",
        "      if self.contextual:\n",
        "        negatives_pred = self(\n",
        "            (negatives_batch['design'], negatives_batch['context_id']),\n",
        "            training=training)\n",
        "      else:\n",
        "        negatives_pred = self(negatives_batch['design'], training=True)\n",
        "    \n",
        "    negatives_pred = tf.clip_by_value(negatives_pred, clip_value_min=-4000.0,\n",
        "                                      clip_value_max=4000.0)\n",
        "    \n",
        "    cql_loss = tf.reduce_mean(negatives_pred)\n",
        "    cql_loss = tf.clip_by_value(cql_loss, \n",
        "                                clip_value_min=-4000, \n",
        "                                clip_value_max=1e6)\n",
        "    loss_dict['negatives_dist'] = tf.reduce_mean(negatives_pred)\n",
        "\n",
        "    mse_loss = weighted_mse_loss(\n",
        "        model_pred, data_batch['objective'], weights)\n",
        "\n",
        "    if loss_type == 'mse+rank':\n",
        "      if self.contextual:\n",
        "        avg_ranking_train_loss = ranking_loss_fn(\n",
        "            model_pred, data_batch['objective'],\n",
        "            context=data_batch['raw_context'])\n",
        "      else:\n",
        "        avg_ranking_train_loss = ranking_loss_fn(\n",
        "            model_pred, data_batch['objective'])\n",
        "    else:\n",
        "      avg_ranking_train_loss = 0.0\n",
        "\n",
        "    # Only used for logging, measures how big the MSE error is relative to\n",
        "    # the output of the model. \n",
        "    avg_approx_loss = weighted_approx_loss(\n",
        "        model_pred, data_batch['objective'], weights)\n",
        "    passed_context = None\n",
        "\n",
        "    if self.contextual:\n",
        "      passed_context = data_batch['raw_context']\n",
        "    \n",
        "    avg_ranking_loss = ranking_loss(\n",
        "        model_pred, data_batch['objective'], context=passed_context)\n",
        "    avg_kendall_loss = kendall_correlation(\n",
        "        model_pred, data_batch['objective'], context=passed_context)\n",
        "\n",
        "    train_loss = mse_loss\n",
        "    loss_dict['mse_loss'] = mse_loss\n",
        "    loss_dict['avg_approx_loss'] = avg_approx_loss\n",
        "    loss_dict['avg_ranking_loss'] = avg_ranking_loss\n",
        "    loss_dict['avg_ranking_train_loss'] = avg_ranking_train_loss\n",
        "    loss_dict['avg_kendall_loss'] = avg_kendall_loss\n",
        "    loss_dict['cql_loss'] = cql_loss\n",
        "    loss_dict['negatives_pred'] = tf.reduce_mean(negatives_pred)\n",
        "    loss_dict['model_pred_average'] = tf.reduce_mean(model_pred)\n",
        "    train_loss = train_loss + ranking_penalty_weight * avg_ranking_train_loss\n",
        "    train_loss = train_loss + self.cql_alpha_value * cql_loss\n",
        "    \n",
        "    if inp_batch_type is not 'valid':\n",
        "      weights_negatives = tf.ones_like(data_batch['objective'])\n",
        "      if self.contextual:\n",
        "        model_pred_invalid, invalid_dict = self(\n",
        "          inputs=(data_batch['invalid/design'], data_batch['context_id']),\n",
        "          training=training, with_logging=True)\n",
        "      else:\n",
        "        model_pred_invalid, invalid_dict = self(\n",
        "            data_batch['invalid/design'], training=training, with_logging=True)\n",
        "\n",
        "      for key in invalid_dict:\n",
        "        loss_dict['invalid/'+key] = invalid_dict[key]\n",
        "\n",
        "      ## Conservatism training\n",
        "      loss_dict['y_value_infeasible'] = tf.reduce_mean(model_pred_invalid)\n",
        "      loss_dict['y_value_infeasible'] = tf.clip_by_value(\n",
        "          loss_dict['y_value_infeasible'], \n",
        "          clip_value_min=-1000, clip_value_max=1e6)\n",
        "      train_loss = train_loss + self.infeasible_alpha *\\\n",
        "            loss_dict['y_value_infeasible']\n",
        "\n",
        "      mse_loss_invalid = weighted_mse_loss(\n",
        "          model_pred_invalid, data_batch['invalid/objective'], \n",
        "          weights_negatives)\n",
        "      avg_approx_loss_invalid = weighted_approx_loss(\n",
        "          model_pred_invalid, data_batch['invalid/objective'], \n",
        "          weights_negatives)\n",
        "      mse_loss = mse_loss + mse_loss_invalid\n",
        "      loss_dict['mse_loss_invalid'] = mse_loss_invalid\n",
        "      loss_dict['mse_loss_overall'] = mse_loss\n",
        "      loss_dict['avg_approx_loss_invalid'] = avg_approx_loss_invalid\n",
        "    return loss_dict, train_loss\n",
        "\n",
        "  def perform_training(self, batch, loss_type,\n",
        "                       ranking_penalty_weight=0.0, **kwargs):\n",
        "    \"\"\"\n",
        "    Actually perform training by computing loss, and then taking gradients\n",
        "    through it. Makes sure to backpropagate through all networks.\n",
        "    \"\"\"\n",
        "    with tf.GradientTape(\n",
        "        watch_accessed_variables=False, persistent=True) as tape:\n",
        "      tape.watch(\n",
        "          [v for net in self.optimize_networks\\\n",
        "           for v in net.trainable_variables])\n",
        "      loss_dict, loss_train = self.compute_loss(\n",
        "          batch, loss_type, training=True,\n",
        "          ranking_penalty_weight=ranking_penalty_weight)\n",
        "\n",
        "    grads = tape.gradient(loss_train,\n",
        "                          [v for net in self.optimize_networks\\\n",
        "                           for v in net.trainable_variables])\n",
        "    gen_grads_op = self.optimizer.apply_gradients(\n",
        "        zip(grads, [v for net in self.optimize_networks\\\n",
        "                    for v in net.trainable_variables]))\n",
        "    return loss_dict\n",
        "\n",
        "  def measure_stats(self, batch, batch_type=None, **kwargs):\n",
        "    \"\"\"Simply make a forward pass through compute_loss to measure losses.\"\"\"\n",
        "    loss_dict, _ = self.compute_loss(batch, loss_type='mse+rank',\n",
        "                                     training=False,\n",
        "                                     inp_batch_type=batch_type)\n",
        "    return loss_dict\n",
        "\n",
        "  def infer_negatives(self, batch):\n",
        "    \"\"\"Run gradient descent to obtain negative examples\"\"\"\n",
        "    temp_batch = dict()\n",
        "    log_probs = batch['design']\n",
        "    if self.contextual:\n",
        "      contexts = batch['context_id']\n",
        "    for _ in range(self.num_gradient_infer_steps):\n",
        "      with tf.GradientTape(\n",
        "          watch_accessed_variables=False, persistent=False) as tape:\n",
        "        tape.watch(log_probs)\n",
        "        if self.contextual:\n",
        "          model_pred = self((log_probs, contexts), training=False)\n",
        "        else:\n",
        "          model_pred = self(log_probs, training=False)\n",
        "      grad = tape.gradient(model_pred, log_probs)\n",
        "      log_probs = log_probs + self.opt_lr * grad[0]\n",
        "    temp_batch['design'] = tf.stop_gradient(log_probs)\n",
        "    if 'context_id' in batch and self.contextual:\n",
        "      temp_batch['context_id'] = batch['context_id']\n",
        "    return temp_batch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1RPktU1SVbIK"
      },
      "source": [
        "# Data Loading and Problem Definition"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uvF6OMaCVlaf"
      },
      "source": [
        "## Hardware Optimization Problem \u0026 Offline Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "mg6CWNYlVkqO"
      },
      "outputs": [],
      "source": [
        "#@title Define the hardware optimization problem\n",
        "\n",
        "class HardwareOptProblem:\n",
        "  \"\"\"\n",
        "  Problem for loading the task dataset and training\n",
        "  \"\"\"\n",
        "  def __init__(self,\n",
        "               config: dict,\n",
        "               data_file: dict, \n",
        "               params_dict: Optional[dict] = None):\n",
        "    \"\"\"Initialize a hardware optimization problem.\n",
        "\n",
        "    config: a dictionary of various input fields and their corresponding\n",
        "      possible valid number of discrete values. \n",
        "    data_file: a dictionary of a list of various input fields.\n",
        "    params_dict: a dictionary of additional inputs to the HardwareOptProblem.\n",
        "    \"\"\"\n",
        "\n",
        "    # Batch size for the batch sampling\n",
        "    self.batch_size = 256\n",
        "    if 'batch_size' in params_dict:\n",
        "      self._batch_size = params_dict['batch_size']\n",
        "\n",
        "    # Whether to train on infeasible points or not\n",
        "    # use 'valid' for feasible points, and 'mixed' for both infeasible and\n",
        "    # feasible points\n",
        "    self._batch_type = 'valid'\n",
        "    if 'batch_type' in params_dict:\n",
        "      self._batch_type = params_dict['batch_type']\n",
        "\n",
        "    # Add any area constraints or not: this flag enables filtering the data\n",
        "    # basedn on whether the area constraint is not satisfied\n",
        "    self._add_area_constraints = False\n",
        "    if 'add_area_constraints' in params_dict:\n",
        "      self._add_area_constraints = params_dict['add_area_constraints']\n",
        "\n",
        "    self.dataset = PRIMEDataset(config=config,\n",
        "                                data_dict=data_file)\n",
        "    self.feasible_probs,\\\n",
        "             self.infeasible_probs = self.dataset.get_feasible_probs(\n",
        "                              add_area_constraints=self._add_area_constraints)\n",
        "\n",
        "    # Choose what kind of batch to provide while training the model\n",
        "    self.get_training_batch = None\n",
        "    if self._batch_type == 'valid':\n",
        "      self.get_training_batch = self.get_valid_only_batch\n",
        "    elif self._batch_type == 'mixed':\n",
        "      self.get_training_batch = self.get_mixed_batch\n",
        "    else:\n",
        "      self.get_training_batch = self.get_all_batch\n",
        "\n",
        "  def get_all_batch(self,):\n",
        "    \"\"\"Sample i.i.d. from the entire dataset.\"\"\"\n",
        "    indices = np.random.randint(1, \n",
        "                                self.dataset._top, self._batch_size)\n",
        "    batch_x, batch_y = self.dataset._get_batch(indices)\n",
        "    batch_dict = dict()\n",
        "    batch_dict['design'] = batch_x\n",
        "    batch_dict['objective'] = batch_y\n",
        "    return batch_dict\n",
        "\n",
        "  def get_valid_only_batch(self,):\n",
        "    \"\"\"Get only valid samples in the batch.\"\"\"\n",
        "    indices = np.random.choice(np.arange(0, self.dataset._top),\n",
        "                               size=self._batch_size, p=self.feasible_probs)\n",
        "    batch_x, batch_y = self.dataset._get_batch(indices)\n",
        "    batch_dict = dict()\n",
        "    batch_dict['design'] = batch_x\n",
        "    batch_dict['objective'] = batch_y\n",
        "    return batch_dict\n",
        "\n",
        "  def get_top_batch(self,):\n",
        "    \"\"\"Get only the top scoring batch for eval\"\"\"\n",
        "    indices = self.dataset._tf_dataset['argsort'][-self.batch_size:]\n",
        "    batch_x, batch_y = self.dataset._get_batch(indices)\n",
        "    batch_dict = dict()\n",
        "    batch_dict['design'] = batch_x\n",
        "    batch_dict['objective'] = batch_y\n",
        "    return batch_dict\n",
        "\n",
        "  def get_mixed_batch(self,):\n",
        "    \"\"\"Get both valid and invalid samples to train in a batch\"\"\"\n",
        "    # Should be called when training with invalid samples as negatives\n",
        "    valid_indices = np.random.choice(np.arange(0, self.dataset._top),\n",
        "                                     size=self._batch_size,\n",
        "                                     p=self.feasible_probs)\n",
        "    invalid_indices = np.random.choice(np.arange(0, self.dataset._top),\n",
        "                                       size=self._batch_size,\n",
        "                                       p=self.infeasible_probs)\n",
        "    batch_x, batch_y = self.dataset._get_batch(valid_indices)\n",
        "    batch_x_in, batch_y_in = self.dataset._get_batch(invalid_indices)\n",
        "    batch_dict = dict()\n",
        "    batch_dict['design'] = batch_x\n",
        "    batch_dict['objective'] = batch_y\n",
        "    batch_dict['invalid/design'] = batch_x_in\n",
        "    batch_dict['invalid/objective'] = batch_y_in\n",
        "    return batch_dict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "b1WllMw1K-Ad"
      },
      "outputs": [],
      "source": [
        "#@title Define the dataset\n",
        "\n",
        "class PRIMEDataset(tf.Module):\n",
        "  \"\"\"\n",
        "  Load the dataset to be able to train the PRIMETransformerModel.  \n",
        "  \"\"\"\n",
        "  def __init__(self,\n",
        "               config,\n",
        "               data_dict: dict,\n",
        "               **kwargs):\n",
        "    \"\"\"Create a dataset for training PRIME.\"\"\"\n",
        "    self._config = config\n",
        "\n",
        "    self.data_dict = data_dict\n",
        "    self._design_space_dict = {}\n",
        "    self._segment_lengths = {}\n",
        "    self._max_ctr = 0\n",
        "    self._eval_metric_keys = ['area', 'runtime', 'score']\n",
        "    self._validity_keys = ['infeasible',]\n",
        "\n",
        "    self._active_training_keys = ['param_1', 'param_2', 'param_3',\n",
        "                                  'param_4', 'param_5', 'param_6',\n",
        "                                  'param_7', 'param_8', 'param_9',\n",
        "                                  'param_10']\n",
        "\n",
        "    self._tf_dataset = {}\n",
        "    self._top = 0\n",
        "    if self.data_dict is not None:\n",
        "      self._setup_dataset()\n",
        "\n",
        "  def _setup_dataset(self,):\n",
        "    \"\"\"Main function to setup the dataset\"\"\"\n",
        "    self.load_or_refresh_config()\n",
        "    logging.info('Loading dataset..')\n",
        "    self._convert_to_tf_dataset()\n",
        "    self.get_score_function()\n",
        "    print ('Loaded dataset....', self.size)\n",
        "\n",
        "  def get_input_splits(self,):\n",
        "    \"\"\"Get the splits of input of the dataset.\"\"\"\n",
        "    lengths = []\n",
        "    for key in self._active_training_keys:\n",
        "      ctr_idx = self._design_space_dict[key]['ctr']\n",
        "      lengths.append(self._segment_lengths[ctr_idx])\n",
        "    self._active_lengths = lengths\n",
        "    return lengths\n",
        "\n",
        "  def get_score_function(self,):\n",
        "    \"\"\"Get the objective function which is being maximized\"\"\"\n",
        "    runtime = self._tf_dataset['runtime'].numpy()\n",
        "    area = self._tf_dataset['area'].numpy()\n",
        "    scores = -runtime\n",
        "    self._tf_dataset['score'] = tf.convert_to_tensor(\n",
        "        scores, dtype=tf.float32)\n",
        "    print ('Score stats: ')\n",
        "    print ('--------------------------------------------')\n",
        "    print ('Max: ', scores.max())\n",
        "    print ('Mean: ', scores.mean())\n",
        "    print ('Min: ', scores.min())\n",
        "    print ('--------------------------------------------')\n",
        "\n",
        "    # Since we need top batch for eval, store top scores\n",
        "    self._tf_dataset['argsort'] = np.argsort(\n",
        "        self._tf_dataset['score'].numpy())\n",
        "    return scores\n",
        "\n",
        "  def _convert_to_tf_dataset(self,):\n",
        "    \"\"\"Convert the dataset to a tensorflow dataset, easy to read from.\"\"\"\n",
        "    tf_dataset = {}\n",
        "    for key in self._active_training_keys +\\\n",
        "        self._eval_metric_keys + self._validity_keys:\n",
        "      tf_dataset[key] = []\n",
        "\n",
        "    # Load the data from the data file. Note that most of the fields are\n",
        "    # actually not one-hots, and essentially corresponds to the original data\n",
        "    # with field-value pairs for each field, and the value is a discrete value.\n",
        "    tf_actual_dataset = {}\n",
        "    parsed_dataset = self.data_dict\n",
        "    for p in parsed_dataset:\n",
        "      tf_dataset[p] = parsed_dataset[p]\n",
        "      tf_actual_dataset[p] =  tf.convert_to_tensor(tf_dataset[p])\n",
        "\n",
        "      if key in self._active_training_keys:\n",
        "        tf_actual_dataset[p] = tf.cast(tf_actual_dataset[p], tf.int32)\n",
        "\n",
        "    self._design_space_dict_copy = deepcopy(self._design_space_dict)\n",
        "\n",
        "    # Now convert the dataset to actually use one-hot representations. This is\n",
        "    # used for training, and so it is important to use this.\n",
        "    tf_actual_temp_dataset = {}\n",
        "    for key in self._active_training_keys:\n",
        "      design_space_map = dict(\n",
        "          self._design_space_dict[key]['mapping_one_hot_to_value'])\n",
        "      data_val = tf_actual_dataset[key].numpy().astype(np.int32).tolist()\n",
        "      out_vals = []\n",
        "      for x in data_val:\n",
        "        out_vals.append(design_space_map[x])\n",
        "\n",
        "      tf_actual_temp_dataset[key] = tf.constant(out_vals, dtype=tf.int32)\n",
        "\n",
        "    ## Finally load the tf_actual_temp_dataset into the tf_dataset\n",
        "    for key in tf_actual_temp_dataset:\n",
        "      tf_actual_dataset[key] = tf_actual_temp_dataset[key]\n",
        "\n",
        "    self._tf_dataset = tf_actual_dataset\n",
        "    self._infeasible_np = self._tf_dataset['infeasible'].numpy().astype(\n",
        "        np.float32)\n",
        "    self._top = self._infeasible_np.shape[0]\n",
        "\n",
        "  def load_or_refresh_config(self):\n",
        "    \"\"\"Load config file with specifications.\"\"\"\n",
        "    self._design_space_dict = {}\n",
        "    self._segment_lengths = {}\n",
        "    \n",
        "    try:\n",
        "      # The case when the config is a file to open\n",
        "      with gfile.Open(self._config, 'r') as f:\n",
        "        line = f.readline()\n",
        "        line = line.replace('\\n', '')\n",
        "        # print ('Line: ', line)\n",
        "        ctr = 0\n",
        "        while line:\n",
        "          ind_field = dict()\n",
        "          split_line = line.split(':')\n",
        "          ind_field['data_type'] = split_line[0]\n",
        "          ind_field['value_range'] = [int(x) for x in split_line[-1].split(',')]\n",
        "          index_vals = np.arange(len(ind_field['value_range']))\n",
        "          ind_field['mapping_one_hot_to_value'] = zip(\n",
        "              ind_field['value_range'], index_vals)\n",
        "          ind_field['ctr'] = ctr\n",
        "          self._design_space_dict[split_line[1]] = ind_field\n",
        "          self._segment_lengths[ctr] = len(ind_field['value_range'])\n",
        "          self._max_ctr += 1\n",
        "          line = f.readline()\n",
        "          ctr += 1\n",
        "    except:\n",
        "      # When config is a string of the contents of the file\n",
        "      lines = self._config.split(\"\\n\")\n",
        "      lines = [line.replace('\\n', '') for line in lines]\n",
        "      ctr = 0\n",
        "      for line in lines:\n",
        "        ind_field = dict()\n",
        "        split_line = line.split(':')\n",
        "        ind_field['data_type'] = split_line[0]\n",
        "        ind_field['value_range'] = [int(x) for x in split_line[-1].split(',')]\n",
        "        index_vals = np.arange(len(ind_field['value_range']))\n",
        "        ind_field['mapping_one_hot_to_value'] = zip(\n",
        "            ind_field['value_range'], index_vals)\n",
        "        ind_field['ctr'] = ctr\n",
        "        self._design_space_dict[split_line[1]] = ind_field\n",
        "        self._segment_lengths[ctr] = len(ind_field['value_range'])\n",
        "        self._max_ctr += 1\n",
        "        ctr += 1\n",
        "\n",
        "    split_lengths = []\n",
        "    for key in self._active_training_keys:\n",
        "      split_lengths.append(\n",
        "          self._segment_lengths[self._design_space_dict[key]['ctr']])\n",
        "    total_length_split = 0\n",
        "    \n",
        "    if total_length_split \u003e 0:\n",
        "      split_lengths.append(total_length_split)\n",
        "    self.split_lengths = split_lengths  # later used to split input when needed\n",
        "    self.continuous_or_not = (total_length_split \u003e 0)\n",
        "\n",
        "  @property\n",
        "  def size(self,):\n",
        "    return self._top\n",
        "\n",
        "  @property\n",
        "  def input_properties(self):\n",
        "    \"\"\"Get the total length of the vector to be fed as input to the model.\"\"\"\n",
        "    length = 0\n",
        "    for val in self._active_lengths:\n",
        "      length += val\n",
        "    return length\n",
        "\n",
        "  def get_feasible_probs(self, add_area_constraints=False):\n",
        "    \"\"\"\n",
        "    Get the probability of points that are feasible, meaning they don't\n",
        "    violate the area constraint and also obtain the feasibility result.  \n",
        "    \"\"\"\n",
        "    feasible = (1.0 - self._infeasible_np)\n",
        "    print ('Number of feasible points: ', np.sum(feasible))\n",
        "    if add_area_constraints:\n",
        "      print ('Min area: ', tf.reduce_min(self._tf_dataset['area']))\n",
        "      feasible_area = (\n",
        "          self._tf_dataset['area'] \u003c= AREA_THRESHOLD).numpy().astype(np.float32)\n",
        "      feasible = np.clip(feasible + feasible_area - 1.0,\n",
        "                         a_min=0.0, a_max=1.0)\n",
        "      print ('Number of feasible points due to area constraint: ',\n",
        "             np.sum(feasible_area))\n",
        "      print ('NUmber of feasible points after area constraint: ',\n",
        "             np.sum(feasible))\n",
        "    probs = feasible / np.sum(feasible)\n",
        "    infeasible_probs = (1.0 - feasible)/ np.sum(1.0 - feasible)\n",
        "    return probs, infeasible_probs\n",
        "\n",
        "  def valid_invalid_data_size(self, add_area_constraints=True):\n",
        "    \"\"\"Get the size of the valid and invalid dataset compositions.\"\"\"\n",
        "    feasible = (1.0 - self._infeasible_np)\n",
        "    if add_area_constraints:\n",
        "      feasible_area = (\n",
        "          self._tf_dataset['area'] \u003c= AREA_THRESHOLD).numpy().astype(np.float32)\n",
        "      feasible = np.clip(feasible + feasible_area - 1.0,\n",
        "                         a_min=0.0, a_max=1.0)\n",
        "    return np.sum(feasible), np.shape(feasible)[0] - np.sum(feasible)\n",
        "\n",
        "  def _get_batch(self, indices):\n",
        "    \"\"\"Sample a batch from the dataset.\"\"\"\n",
        "    all_train_elements = []  # this is the training elements in one-hot form\n",
        "    all_test_elements = []  # this is the evaluation fields (area, runtime, score, etc)\n",
        "\n",
        "    # Discrete training input keys\n",
        "    for key in self._active_training_keys:\n",
        "      all_train_elements.append(\n",
        "              tf.one_hot(tf.gather(self._tf_dataset[key], indices),\n",
        "                         depth=self._segment_lengths[\n",
        "                                self._design_space_dict[key]['ctr']]))\n",
        "\n",
        "    # Eval keys\n",
        "    all_test_elements = tf.expand_dims(\n",
        "        tf.gather(self._tf_dataset['score'], indices), 1)\n",
        "    return tf.concat(all_train_elements, 1), all_test_elements"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X3XTzmx2XEmu"
      },
      "source": [
        "# Training loop and training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "KZt31sS_Xcuz"
      },
      "outputs": [],
      "source": [
        "#@title Defining the function that runs training\n",
        "\n",
        "def train_eval_offline(\n",
        "    # Data flags\n",
        "    config=None,\n",
        "    training_dataset=None,\n",
        "    validation_dataset=None,    \n",
        "    # Train flags\n",
        "    train_steps=int(1e6),\n",
        "    summary_freq=1000,\n",
        "    eval_freq=1000,\n",
        "    # Train hparams\n",
        "    add_summary=True,\n",
        "    save_dir=None,\n",
        "    loss_type='mse',\n",
        "    layers=(512, 512, 512),\n",
        "    opt_lr=1e-4,\n",
        "    opt_betas=(0.9, 0.999),\n",
        "    with_ranking_penalty=False,\n",
        "    ranking_penalty_weight=0.1,\n",
        "    batch_size=256,\n",
        "    batch_type='mixed',\n",
        "    # params of the model\n",
        "    use_dropout=False,\n",
        "    num_votes=1,\n",
        "    # PRIME parameters:\n",
        "    cql_alpha=1.0,\n",
        "    infeasible_alpha=1.0):\n",
        "  \"\"\"Training loop for the PRIME model. \n",
        "  \n",
        "  Most of the input arguments are primarily hyperparameters for training the\n",
        "  PRIME model, and self explanatory. Other arguments explained below. \n",
        "  \n",
        "  save_dir: the directory where the store the saved model, and the training\n",
        "    summaries. Can be a string or None.\n",
        "  training_dataset: a dictionary of fields in the training dataset, and their\n",
        "    corresponding values used to train.\n",
        "  validation_dataset: a dictionary of fields in the validation dataset, and\n",
        "    their corresponding values to measure cross-validation. \n",
        "  \"\"\"\n",
        "\n",
        "  # First create the training dataset, note that the dataset below is a\n",
        "  # dummy dataset, that is only well-suited for training as a representative\n",
        "  # example. You can plug in the dataset from the other colab that provides\n",
        "  # the data for training, or you can add your own dataset here.  \n",
        "  params_dict = dict()\n",
        "  params_dict['batch_size'] = batch_size\n",
        "  params_dict['batch_type'] = batch_type\n",
        "  params_dict['add_area_constraints'] = True\n",
        "  # Defining the problem automatically does dataset loading\n",
        "  train_problem = HardwareOptProblem(config,\n",
        "                                    training_dataset, params_dict)\n",
        "  \n",
        "  # Now define the validation dataset (or val_problem)\n",
        "  val_params_dict = dict()\n",
        "  val_params_dict['batch_size'] = batch_size\n",
        "  val_params_dict['add_area_constraints'] = True\n",
        "  # Only validate on the valid samples in the validation dataset\n",
        "  val_params_dict['batch_type'] = 'valid'\n",
        "  val_problem = HardwareOptProblem(config, validation_dataset,\n",
        "                                   val_params_dict)\n",
        "\n",
        "\n",
        "  # The dimensionality of each parameter. this input_splits parameter goes\n",
        "  # into the PRIMETransformer, as it enables us to pass in inputd as a big\n",
        "  # vector of concatenated one-hot vectors for each discrete parameter, and\n",
        "  # then unpack it in the model training. This gives the flexibility of actually\n",
        "  # being able to use the input one-hot vectors in any way as needed. \n",
        "  input_splits = train_problem.dataset.get_input_splits()\n",
        "  print ('Input splits: ', input_splits)\n",
        "\n",
        "  # Number of inputs in all: the total dimensionality of the input is given by\n",
        "  # the sum of number of possible values each discrete parameter can take\n",
        "  input_properties = train_problem.dataset.input_properties\n",
        "  print ('Loaded validation dataset..', train_problem.dataset.size, \n",
        "          val_problem.dataset.size, input_properties)\n",
        "\n",
        "  feasible_size,\\\n",
        "    infeasible_size = train_problem.dataset.valid_invalid_data_size()\n",
        "  print ('Feasible/Infeasible size: ', feasible_size, infeasible_size)\n",
        "\n",
        "  fwd_optimizer = tf.keras.optimizers.Adam(learning_rate=opt_lr,\n",
        "                                           beta_1=opt_betas[0],\n",
        "                                           beta_2=opt_betas[1], name='opt')\n",
        "\n",
        "  training_dict = dict()\n",
        "  training_dict['training_type'] = batch_type\n",
        "  training_dict['use_dropout'] = use_dropout\n",
        "  training_dict['infeasible_alpha'] = infeasible_alpha\n",
        "  training_dict['input_splits'] = input_splits\n",
        "  training_dict['num_votes'] = num_votes\n",
        "  training_dict['infeasbile_multiplier'] = float(feasible_size)/(\n",
        "      float(infeasible_size) + 1)\n",
        "  training_dict['num_gradient_steps'] = 20\n",
        "\n",
        "  model = PRIMETransformerModel(\n",
        "        num_outputs=1,\n",
        "        num_inputs=input_properties,\n",
        "        optimizer=fwd_optimizer,\n",
        "        layers=layers,\n",
        "        penalty_weight=cql_alpha,\n",
        "        params_dict=training_dict)\n",
        "  \n",
        "  rand_num = np.random.randint(10000)\n",
        "  \n",
        "  # summary writer\n",
        "  if save_dir is not None:\n",
        "    save_dir = os.path.join(save_dir, str(rand_num))\n",
        "    summary_writer = tf.summary.create_file_writer(logdir=save_dir)\n",
        "    summary_writer.set_as_default()\n",
        "  else:\n",
        "    tf.summary.create_noop_writer()\n",
        "\n",
        "  print ('save dir : ', save_dir)\n",
        "\n",
        "  # Now start the training\n",
        "  for step in range(train_steps):\n",
        "    batch = train_problem.get_training_batch()\n",
        "    # This is just to build the models.\n",
        "    if step == 0:\n",
        "      _ = model.measure_stats(batch)\n",
        "    loss_dict = model.perform_training(\n",
        "        batch, loss_type=loss_type,\n",
        "        ranking_penalty_weight=ranking_penalty_weight)\n",
        "\n",
        "    if step % summary_freq == 0:\n",
        "      # regular logging\n",
        "      print ('-------------------------------------------------------')\n",
        "      for key in loss_dict:\n",
        "        tf.summary.scalar('train/' + key, loss_dict[key], step=step)\n",
        "        print ('Step: ', step, 'train/' + key, ':', loss_dict[key])\n",
        "      print ('-------------------------------------------------------')\n",
        "\n",
        "      if save_dir is not None:\n",
        "        if step == 0:\n",
        "          model.save(save_dir)\n",
        "        if step % 5000 == 0:\n",
        "          model.save_weights(os.path.join(save_dir, \"ckpt-\"+str(step)))\n",
        "\n",
        "    if step % eval_freq == 0:\n",
        "      val_batch = val_problem.get_training_batch()\n",
        "      # validation batches are only valid batches\n",
        "      val_loss_dict = model.measure_stats(val_batch, batch_type='valid')\n",
        "      print ('-------------------------------------------------------')\n",
        "      for key in val_loss_dict:\n",
        "        tf.summary.scalar('val/' + key, val_loss_dict[key], step=step)\n",
        "        print ('Step: ', step, 'val/' + key, ':', val_loss_dict[key])\n",
        "      print ('-------------------------------------------------------')\n",
        "\n",
        "  print ('Finished Training')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sjN3ZcR3kLnt",
        "outputId": "b92be4a6-c5e2-42a1-c5c2-e5b1ba31dc51"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Keys in the dummy dataset:  dict_keys(['param_1', 'param_2', 'param_3', 'param_7', 'param_8', 'param_9', 'param_6', 'param_5', 'param_4', 'param_10', 'area', 'runtime', 'infeasible'])\n"
          ]
        }
      ],
      "source": [
        "#@title Creating dummy data\n",
        "\n",
        "\n",
        "def create_dummy_data(dataset_size, config_dict,\n",
        "                      remaining_fields=('area', 'runtime', 'infeasible')):\n",
        "  \"\"\"\n",
        "  Create dummy data for training PRIME.\n",
        "\n",
        "  dataset_size: int, the number of samples in the training data.\n",
        "  config_dict: dictionary, specifying the design space to sample the dataset.\n",
        "  remaining_fields: tuple, a tuple of float values specifying additional fields\n",
        "    to sample in the dataset. Default includes float parameters 'area' and\n",
        "    'runtime', and also includes a boolean field 'infeasible'\n",
        "  \"\"\"\n",
        "  data_dict = {}\n",
        "  for key in config_dict:\n",
        "    allowed_values = config_dict[key]['value_range']\n",
        "    sampled_values = np.random.choice(allowed_values, size=dataset_size)\n",
        "    data_dict[key] = sampled_values\n",
        "\n",
        "  for key in remaining_fields:\n",
        "    if 'infeasible' in key:\n",
        "      sampled_values = np.random.choice([True, False], size=dataset_size)\n",
        "    elif 'runtime' or 'area' in key:\n",
        "      sampled_values = np.random.uniform(low=0.0, high=100.0,\n",
        "                                         size=dataset_size)\n",
        "    data_dict[key] = sampled_values\n",
        "\n",
        "  return data_dict\n",
        "\n",
        "\n",
        "config_str = \"\"\"discrete:param_1:float64:true:1,2,4,6,8,10,12,14,16,32\n",
        "discrete:param_2:float64:true:1,2,4,6,8,10,12,14,16,32\n",
        "discrete:param_3:float64:true:4,8,16,32,64,128,256\n",
        "discrete:param_7:float64:true:256,512,1024,2048,4096,8192,16384\n",
        "discrete:param_8:float64:true:8192,16384,32768,65536\n",
        "discrete:param_9:float64:true:2048,4096,8192,16384,32768\n",
        "discrete:param_6:float64:true:4096,8192,16384,32768,65536,131072,262144,524288,1048576,2097152,4194304\n",
        "discrete:param_5:float64:true:262144,524288,1048576,2097152,4194304,8388608,16777216\n",
        "discrete:param_4:float64:true:1,2,4,6,8,10,12,14,16,32\n",
        "discrete:param_10:float64:true:5,10,16,20,25,30\"\"\"\n",
        "\n",
        "# Generating dummy data\n",
        "temp_dataset = PRIMEDataset(config=config_str, data_dict=None, params_dict={})\n",
        "temp_dataset.load_or_refresh_config()\n",
        "dictionary_to_generate_data = temp_dataset._design_space_dict\n",
        "\n",
        "training_data = create_dummy_data(1000, dictionary_to_generate_data)\n",
        "validation_data = create_dummy_data(1000, dictionary_to_generate_data)\n",
        "\n",
        "print ('Keys in the dummy dataset: ', training_data.keys())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_V1oRa25sRYV",
        "outputId": "2dfc8315-1abd-47f1-f12e-0a6072bc99f7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Score stats: \n",
            "--------------------------------------------\n",
            "Max:  -0.6129692603849213\n",
            "Mean:  -51.45989737306953\n",
            "Min:  -99.96909843043453\n",
            "--------------------------------------------\n",
            "Loaded dataset.... 1000\n",
            "Number of feasible points:  511.0\n",
            "Min area:  tf.Tensor(0.06833098587071262, shape=(), dtype=float64)\n",
            "Number of feasible points due to area constraint:  295.0\n",
            "NUmber of feasible points after area constraint:  153.0\n",
            "Score stats: \n",
            "--------------------------------------------\n",
            "Max:  -0.04601063577058806\n",
            "Mean:  -49.507973451263794\n",
            "Min:  -99.947081272807\n",
            "--------------------------------------------\n",
            "Loaded dataset.... 1000\n",
            "Number of feasible points:  500.0\n",
            "Min area:  tf.Tensor(0.037564213198348906, shape=(), dtype=float64)\n",
            "Number of feasible points due to area constraint:  255.0\n",
            "NUmber of feasible points after area constraint:  123.0\n",
            "Input splits:  [10, 10, 7, 10, 7, 11, 7, 4, 5, 6]\n",
            "Loaded validation dataset.. 1000 1000 77\n",
            "Feasible/Infeasible size:  153.0 847.0\n",
            "Infeasible alpha:  0.05\n",
            "CQL Alpha:  \u003ctf.Variable 'Variable:0' shape=() dtype=float32, numpy=-2.302575\u003e\n",
            "Num votes:  1\n",
            "ListWrapper([10, 10, 7, 10, 7, 11, 7, 4, 5, 6])\n",
            "All networks:  3\n",
            "save dir :  None\n",
            "-------------------------------------------------------\n",
            "Step:  0 train/y_values_max : tf.Tensor(-0.9038583, shape=(), dtype=float32)\n",
            "Step:  0 train/y_values_mean : tf.Tensor(-51.291016, shape=(), dtype=float32)\n",
            "Step:  0 train/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  0 train/negatives_dist : tf.Tensor(0.731879, shape=(), dtype=float32)\n",
            "Step:  0 train/mse_loss : tf.Tensor(3536.3372, shape=(), dtype=float32)\n",
            "Step:  0 train/avg_approx_loss : tf.Tensor(1.0531151, shape=(), dtype=float32)\n",
            "Step:  0 train/avg_ranking_loss : tf.Tensor(0.030993937, shape=(), dtype=float32)\n",
            "Step:  0 train/avg_ranking_train_loss : tf.Tensor(16.001637, shape=(), dtype=float32)\n",
            "Step:  0 train/avg_kendall_loss : tf.Tensor(0.027144551, shape=(), dtype=float32)\n",
            "Step:  0 train/cql_loss : tf.Tensor(0.731879, shape=(), dtype=float32)\n",
            "Step:  0 train/negatives_pred : tf.Tensor(0.731879, shape=(), dtype=float32)\n",
            "Step:  0 train/model_pred_average : tf.Tensor(0.70406383, shape=(), dtype=float32)\n",
            "Step:  0 train/invalid/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  0 train/y_value_infeasible : tf.Tensor(0.7113452, shape=(), dtype=float32)\n",
            "Step:  0 train/mse_loss_invalid : tf.Tensor(3873.5989, shape=(), dtype=float32)\n",
            "Step:  0 train/mse_loss_overall : tf.Tensor(7409.936, shape=(), dtype=float32)\n",
            "Step:  0 train/avg_approx_loss_invalid : tf.Tensor(1.0299487, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  0 val/y_values_max : tf.Tensor(-0.80804974, shape=(), dtype=float32)\n",
            "Step:  0 val/y_values_mean : tf.Tensor(-50.549976, shape=(), dtype=float32)\n",
            "Step:  0 val/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  0 val/negatives_dist : tf.Tensor(0.089412, shape=(), dtype=float32)\n",
            "Step:  0 val/mse_loss : tf.Tensor(3385.7322, shape=(), dtype=float32)\n",
            "Step:  0 val/avg_approx_loss : tf.Tensor(1.0002624, shape=(), dtype=float32)\n",
            "Step:  0 val/avg_ranking_loss : tf.Tensor(-0.014705883, shape=(), dtype=float32)\n",
            "Step:  0 val/avg_ranking_train_loss : tf.Tensor(16.795597, shape=(), dtype=float32)\n",
            "Step:  0 val/avg_kendall_loss : tf.Tensor(0.0042892694, shape=(), dtype=float32)\n",
            "Step:  0 val/cql_loss : tf.Tensor(0.089412, shape=(), dtype=float32)\n",
            "Step:  0 val/negatives_pred : tf.Tensor(0.089412, shape=(), dtype=float32)\n",
            "Step:  0 val/model_pred_average : tf.Tensor(0.017796617, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  10 train/y_values_max : tf.Tensor(-0.9038583, shape=(), dtype=float32)\n",
            "Step:  10 train/y_values_mean : tf.Tensor(-49.595383, shape=(), dtype=float32)\n",
            "Step:  10 train/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  10 train/negatives_dist : tf.Tensor(-5.4266176, shape=(), dtype=float32)\n",
            "Step:  10 train/mse_loss : tf.Tensor(2678.9875, shape=(), dtype=float32)\n",
            "Step:  10 train/avg_approx_loss : tf.Tensor(0.9219241, shape=(), dtype=float32)\n",
            "Step:  10 train/avg_ranking_loss : tf.Tensor(0.041989733, shape=(), dtype=float32)\n",
            "Step:  10 train/avg_ranking_train_loss : tf.Tensor(14.9530525, shape=(), dtype=float32)\n",
            "Step:  10 train/avg_kendall_loss : tf.Tensor(0.03296566, shape=(), dtype=float32)\n",
            "Step:  10 train/cql_loss : tf.Tensor(-5.4266176, shape=(), dtype=float32)\n",
            "Step:  10 train/negatives_pred : tf.Tensor(-5.4266176, shape=(), dtype=float32)\n",
            "Step:  10 train/model_pred_average : tf.Tensor(-5.489807, shape=(), dtype=float32)\n",
            "Step:  10 train/invalid/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  10 train/y_value_infeasible : tf.Tensor(-5.4988065, shape=(), dtype=float32)\n",
            "Step:  10 train/mse_loss_invalid : tf.Tensor(2921.0833, shape=(), dtype=float32)\n",
            "Step:  10 train/mse_loss_overall : tf.Tensor(5600.071, shape=(), dtype=float32)\n",
            "Step:  10 train/avg_approx_loss_invalid : tf.Tensor(0.8700904, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  10 val/y_values_max : tf.Tensor(-0.80804974, shape=(), dtype=float32)\n",
            "Step:  10 val/y_values_mean : tf.Tensor(-48.882214, shape=(), dtype=float32)\n",
            "Step:  10 val/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  10 val/negatives_dist : tf.Tensor(-6.0133524, shape=(), dtype=float32)\n",
            "Step:  10 val/mse_loss : tf.Tensor(2728.4995, shape=(), dtype=float32)\n",
            "Step:  10 val/avg_approx_loss : tf.Tensor(0.8313758, shape=(), dtype=float32)\n",
            "Step:  10 val/avg_ranking_loss : tf.Tensor(-0.25672352, shape=(), dtype=float32)\n",
            "Step:  10 val/avg_ranking_train_loss : tf.Tensor(21.984526, shape=(), dtype=float32)\n",
            "Step:  10 val/avg_kendall_loss : tf.Tensor(-0.1670956, shape=(), dtype=float32)\n",
            "Step:  10 val/cql_loss : tf.Tensor(-6.0133524, shape=(), dtype=float32)\n",
            "Step:  10 val/negatives_pred : tf.Tensor(-6.0133524, shape=(), dtype=float32)\n",
            "Step:  10 val/model_pred_average : tf.Tensor(-6.4133196, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  20 train/y_values_max : tf.Tensor(-0.9038583, shape=(), dtype=float32)\n",
            "Step:  20 train/y_values_mean : tf.Tensor(-52.60445, shape=(), dtype=float32)\n",
            "Step:  20 train/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  20 train/negatives_dist : tf.Tensor(-11.829678, shape=(), dtype=float32)\n",
            "Step:  20 train/mse_loss : tf.Tensor(2489.3354, shape=(), dtype=float32)\n",
            "Step:  20 train/avg_approx_loss : tf.Tensor(0.967807, shape=(), dtype=float32)\n",
            "Step:  20 train/avg_ranking_loss : tf.Tensor(0.009140394, shape=(), dtype=float32)\n",
            "Step:  20 train/avg_ranking_train_loss : tf.Tensor(16.383146, shape=(), dtype=float32)\n",
            "Step:  20 train/avg_kendall_loss : tf.Tensor(0.013235331, shape=(), dtype=float32)\n",
            "Step:  20 train/cql_loss : tf.Tensor(-11.829678, shape=(), dtype=float32)\n",
            "Step:  20 train/negatives_pred : tf.Tensor(-11.829678, shape=(), dtype=float32)\n",
            "Step:  20 train/model_pred_average : tf.Tensor(-11.771761, shape=(), dtype=float32)\n",
            "Step:  20 train/invalid/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  20 train/y_value_infeasible : tf.Tensor(-11.7452755, shape=(), dtype=float32)\n",
            "Step:  20 train/mse_loss_invalid : tf.Tensor(2442.1533, shape=(), dtype=float32)\n",
            "Step:  20 train/mse_loss_overall : tf.Tensor(4931.489, shape=(), dtype=float32)\n",
            "Step:  20 train/avg_approx_loss_invalid : tf.Tensor(0.91133654, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  20 val/y_values_max : tf.Tensor(-0.80804974, shape=(), dtype=float32)\n",
            "Step:  20 val/y_values_mean : tf.Tensor(-51.03371, shape=(), dtype=float32)\n",
            "Step:  20 val/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  20 val/negatives_dist : tf.Tensor(-12.448181, shape=(), dtype=float32)\n",
            "Step:  20 val/mse_loss : tf.Tensor(2303.6575, shape=(), dtype=float32)\n",
            "Step:  20 val/avg_approx_loss : tf.Tensor(0.93867505, shape=(), dtype=float32)\n",
            "Step:  20 val/avg_ranking_loss : tf.Tensor(-0.40264937, shape=(), dtype=float32)\n",
            "Step:  20 val/avg_ranking_train_loss : tf.Tensor(23.87465, shape=(), dtype=float32)\n",
            "Step:  20 val/avg_kendall_loss : tf.Tensor(-0.2642157, shape=(), dtype=float32)\n",
            "Step:  20 val/cql_loss : tf.Tensor(-12.448181, shape=(), dtype=float32)\n",
            "Step:  20 val/negatives_pred : tf.Tensor(-12.448181, shape=(), dtype=float32)\n",
            "Step:  20 val/model_pred_average : tf.Tensor(-13.260455, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  30 train/y_values_max : tf.Tensor(-0.9038583, shape=(), dtype=float32)\n",
            "Step:  30 train/y_values_mean : tf.Tensor(-50.5356, shape=(), dtype=float32)\n",
            "Step:  30 train/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  30 train/negatives_dist : tf.Tensor(-20.171307, shape=(), dtype=float32)\n",
            "Step:  30 train/mse_loss : tf.Tensor(1624.793, shape=(), dtype=float32)\n",
            "Step:  30 train/avg_approx_loss : tf.Tensor(1.0920005, shape=(), dtype=float32)\n",
            "Step:  30 train/avg_ranking_loss : tf.Tensor(-0.12022346, shape=(), dtype=float32)\n",
            "Step:  30 train/avg_ranking_train_loss : tf.Tensor(16.979572, shape=(), dtype=float32)\n",
            "Step:  30 train/avg_kendall_loss : tf.Tensor(-0.07285541, shape=(), dtype=float32)\n",
            "Step:  30 train/cql_loss : tf.Tensor(-20.171307, shape=(), dtype=float32)\n",
            "Step:  30 train/negatives_pred : tf.Tensor(-20.171307, shape=(), dtype=float32)\n",
            "Step:  30 train/model_pred_average : tf.Tensor(-20.278915, shape=(), dtype=float32)\n",
            "Step:  30 train/invalid/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  30 train/y_value_infeasible : tf.Tensor(-20.135735, shape=(), dtype=float32)\n",
            "Step:  30 train/mse_loss_invalid : tf.Tensor(1957.1533, shape=(), dtype=float32)\n",
            "Step:  30 train/mse_loss_overall : tf.Tensor(3581.9463, shape=(), dtype=float32)\n",
            "Step:  30 train/avg_approx_loss_invalid : tf.Tensor(1.0983152, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  30 val/y_values_max : tf.Tensor(-0.80804974, shape=(), dtype=float32)\n",
            "Step:  30 val/y_values_mean : tf.Tensor(-55.31862, shape=(), dtype=float32)\n",
            "Step:  30 val/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  30 val/negatives_dist : tf.Tensor(-21.205795, shape=(), dtype=float32)\n",
            "Step:  30 val/mse_loss : tf.Tensor(1850.8827, shape=(), dtype=float32)\n",
            "Step:  30 val/avg_approx_loss : tf.Tensor(1.027076, shape=(), dtype=float32)\n",
            "Step:  30 val/avg_ranking_loss : tf.Tensor(-0.32453153, shape=(), dtype=float32)\n",
            "Step:  30 val/avg_ranking_train_loss : tf.Tensor(21.168787, shape=(), dtype=float32)\n",
            "Step:  30 val/avg_kendall_loss : tf.Tensor(-0.20484066, shape=(), dtype=float32)\n",
            "Step:  30 val/cql_loss : tf.Tensor(-21.205795, shape=(), dtype=float32)\n",
            "Step:  30 val/negatives_pred : tf.Tensor(-21.205795, shape=(), dtype=float32)\n",
            "Step:  30 val/model_pred_average : tf.Tensor(-22.451515, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  40 train/y_values_max : tf.Tensor(-2.018335, shape=(), dtype=float32)\n",
            "Step:  40 train/y_values_mean : tf.Tensor(-49.35659, shape=(), dtype=float32)\n",
            "Step:  40 train/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  40 train/negatives_dist : tf.Tensor(-31.2863, shape=(), dtype=float32)\n",
            "Step:  40 train/mse_loss : tf.Tensor(1105.0841, shape=(), dtype=float32)\n",
            "Step:  40 train/avg_approx_loss : tf.Tensor(1.0003604, shape=(), dtype=float32)\n",
            "Step:  40 train/avg_ranking_loss : tf.Tensor(-0.05272517, shape=(), dtype=float32)\n",
            "Step:  40 train/avg_ranking_train_loss : tf.Tensor(16.813736, shape=(), dtype=float32)\n",
            "Step:  40 train/avg_kendall_loss : tf.Tensor(-0.029595613, shape=(), dtype=float32)\n",
            "Step:  40 train/cql_loss : tf.Tensor(-31.2863, shape=(), dtype=float32)\n",
            "Step:  40 train/negatives_pred : tf.Tensor(-31.2863, shape=(), dtype=float32)\n",
            "Step:  40 train/model_pred_average : tf.Tensor(-31.335064, shape=(), dtype=float32)\n",
            "Step:  40 train/invalid/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  40 train/y_value_infeasible : tf.Tensor(-31.414925, shape=(), dtype=float32)\n",
            "Step:  40 train/mse_loss_invalid : tf.Tensor(1292.1438, shape=(), dtype=float32)\n",
            "Step:  40 train/mse_loss_overall : tf.Tensor(2397.228, shape=(), dtype=float32)\n",
            "Step:  40 train/avg_approx_loss_invalid : tf.Tensor(1.2672777, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  40 val/y_values_max : tf.Tensor(-0.80804974, shape=(), dtype=float32)\n",
            "Step:  40 val/y_values_mean : tf.Tensor(-47.77986, shape=(), dtype=float32)\n",
            "Step:  40 val/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  40 val/negatives_dist : tf.Tensor(-32.540966, shape=(), dtype=float32)\n",
            "Step:  40 val/mse_loss : tf.Tensor(1008.05023, shape=(), dtype=float32)\n",
            "Step:  40 val/avg_approx_loss : tf.Tensor(1.0481693, shape=(), dtype=float32)\n",
            "Step:  40 val/avg_ranking_loss : tf.Tensor(-0.2252382, shape=(), dtype=float32)\n",
            "Step:  40 val/avg_ranking_train_loss : tf.Tensor(20.3955, shape=(), dtype=float32)\n",
            "Step:  40 val/avg_kendall_loss : tf.Tensor(-0.1397059, shape=(), dtype=float32)\n",
            "Step:  40 val/cql_loss : tf.Tensor(-32.540966, shape=(), dtype=float32)\n",
            "Step:  40 val/negatives_pred : tf.Tensor(-32.540966, shape=(), dtype=float32)\n",
            "Step:  40 val/model_pred_average : tf.Tensor(-34.131866, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  50 train/y_values_max : tf.Tensor(-0.9038583, shape=(), dtype=float32)\n",
            "Step:  50 train/y_values_mean : tf.Tensor(-48.05492, shape=(), dtype=float32)\n",
            "Step:  50 train/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  50 train/negatives_dist : tf.Tensor(-43.863987, shape=(), dtype=float32)\n",
            "Step:  50 train/mse_loss : tf.Tensor(712.9634, shape=(), dtype=float32)\n",
            "Step:  50 train/avg_approx_loss : tf.Tensor(1.8901846, shape=(), dtype=float32)\n",
            "Step:  50 train/avg_ranking_loss : tf.Tensor(0.08773533, shape=(), dtype=float32)\n",
            "Step:  50 train/avg_ranking_train_loss : tf.Tensor(13.866482, shape=(), dtype=float32)\n",
            "Step:  50 train/avg_kendall_loss : tf.Tensor(0.067892194, shape=(), dtype=float32)\n",
            "Step:  50 train/cql_loss : tf.Tensor(-43.863987, shape=(), dtype=float32)\n",
            "Step:  50 train/negatives_pred : tf.Tensor(-43.863987, shape=(), dtype=float32)\n",
            "Step:  50 train/model_pred_average : tf.Tensor(-43.911087, shape=(), dtype=float32)\n",
            "Step:  50 train/invalid/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  50 train/y_value_infeasible : tf.Tensor(-43.836414, shape=(), dtype=float32)\n",
            "Step:  50 train/mse_loss_invalid : tf.Tensor(924.3006, shape=(), dtype=float32)\n",
            "Step:  50 train/mse_loss_overall : tf.Tensor(1637.2639, shape=(), dtype=float32)\n",
            "Step:  50 train/avg_approx_loss_invalid : tf.Tensor(1.7017825, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  50 val/y_values_max : tf.Tensor(-0.80804974, shape=(), dtype=float32)\n",
            "Step:  50 val/y_values_mean : tf.Tensor(-47.65081, shape=(), dtype=float32)\n",
            "Step:  50 val/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  50 val/negatives_dist : tf.Tensor(-44.91657, shape=(), dtype=float32)\n",
            "Step:  50 val/mse_loss : tf.Tensor(808.4212, shape=(), dtype=float32)\n",
            "Step:  50 val/avg_approx_loss : tf.Tensor(1.5485595, shape=(), dtype=float32)\n",
            "Step:  50 val/avg_ranking_loss : tf.Tensor(-0.17357378, shape=(), dtype=float32)\n",
            "Step:  50 val/avg_ranking_train_loss : tf.Tensor(19.457308, shape=(), dtype=float32)\n",
            "Step:  50 val/avg_kendall_loss : tf.Tensor(-0.113051474, shape=(), dtype=float32)\n",
            "Step:  50 val/cql_loss : tf.Tensor(-44.91657, shape=(), dtype=float32)\n",
            "Step:  50 val/negatives_pred : tf.Tensor(-44.91657, shape=(), dtype=float32)\n",
            "Step:  50 val/model_pred_average : tf.Tensor(-47.15584, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  60 train/y_values_max : tf.Tensor(-0.93947476, shape=(), dtype=float32)\n",
            "Step:  60 train/y_values_mean : tf.Tensor(-49.415802, shape=(), dtype=float32)\n",
            "Step:  60 train/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  60 train/negatives_dist : tf.Tensor(-52.542496, shape=(), dtype=float32)\n",
            "Step:  60 train/mse_loss : tf.Tensor(816.1859, shape=(), dtype=float32)\n",
            "Step:  60 train/avg_approx_loss : tf.Tensor(2.075222, shape=(), dtype=float32)\n",
            "Step:  60 train/avg_ranking_loss : tf.Tensor(0.04479072, shape=(), dtype=float32)\n",
            "Step:  60 train/avg_ranking_train_loss : tf.Tensor(15.618368, shape=(), dtype=float32)\n",
            "Step:  60 train/avg_kendall_loss : tf.Tensor(0.0363971, shape=(), dtype=float32)\n",
            "Step:  60 train/cql_loss : tf.Tensor(-52.542496, shape=(), dtype=float32)\n",
            "Step:  60 train/negatives_pred : tf.Tensor(-52.542496, shape=(), dtype=float32)\n",
            "Step:  60 train/model_pred_average : tf.Tensor(-52.71419, shape=(), dtype=float32)\n",
            "Step:  60 train/invalid/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  60 train/y_value_infeasible : tf.Tensor(-52.58152, shape=(), dtype=float32)\n",
            "Step:  60 train/mse_loss_invalid : tf.Tensor(929.5502, shape=(), dtype=float32)\n",
            "Step:  60 train/mse_loss_overall : tf.Tensor(1745.7361, shape=(), dtype=float32)\n",
            "Step:  60 train/avg_approx_loss_invalid : tf.Tensor(2.5227823, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  60 val/y_values_max : tf.Tensor(-0.80804974, shape=(), dtype=float32)\n",
            "Step:  60 val/y_values_mean : tf.Tensor(-50.25917, shape=(), dtype=float32)\n",
            "Step:  60 val/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  60 val/negatives_dist : tf.Tensor(-53.560272, shape=(), dtype=float32)\n",
            "Step:  60 val/mse_loss : tf.Tensor(931.7635, shape=(), dtype=float32)\n",
            "Step:  60 val/avg_approx_loss : tf.Tensor(2.0073552, shape=(), dtype=float32)\n",
            "Step:  60 val/avg_ranking_loss : tf.Tensor(-0.1795484, shape=(), dtype=float32)\n",
            "Step:  60 val/avg_ranking_train_loss : tf.Tensor(20.243542, shape=(), dtype=float32)\n",
            "Step:  60 val/avg_kendall_loss : tf.Tensor(-0.11832106, shape=(), dtype=float32)\n",
            "Step:  60 val/cql_loss : tf.Tensor(-53.560272, shape=(), dtype=float32)\n",
            "Step:  60 val/negatives_pred : tf.Tensor(-53.560272, shape=(), dtype=float32)\n",
            "Step:  60 val/model_pred_average : tf.Tensor(-55.736107, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  70 train/y_values_max : tf.Tensor(-0.9038583, shape=(), dtype=float32)\n",
            "Step:  70 train/y_values_mean : tf.Tensor(-51.05237, shape=(), dtype=float32)\n",
            "Step:  70 train/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  70 train/negatives_dist : tf.Tensor(-53.52229, shape=(), dtype=float32)\n",
            "Step:  70 train/mse_loss : tf.Tensor(747.3913, shape=(), dtype=float32)\n",
            "Step:  70 train/avg_approx_loss : tf.Tensor(1.7664852, shape=(), dtype=float32)\n",
            "Step:  70 train/avg_ranking_loss : tf.Tensor(0.004579853, shape=(), dtype=float32)\n",
            "Step:  70 train/avg_ranking_train_loss : tf.Tensor(15.55942, shape=(), dtype=float32)\n",
            "Step:  70 train/avg_kendall_loss : tf.Tensor(0.009129882, shape=(), dtype=float32)\n",
            "Step:  70 train/cql_loss : tf.Tensor(-53.52229, shape=(), dtype=float32)\n",
            "Step:  70 train/negatives_pred : tf.Tensor(-53.52229, shape=(), dtype=float32)\n",
            "Step:  70 train/model_pred_average : tf.Tensor(-53.431484, shape=(), dtype=float32)\n",
            "Step:  70 train/invalid/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  70 train/y_value_infeasible : tf.Tensor(-53.6995, shape=(), dtype=float32)\n",
            "Step:  70 train/mse_loss_invalid : tf.Tensor(960.54016, shape=(), dtype=float32)\n",
            "Step:  70 train/mse_loss_overall : tf.Tensor(1707.9314, shape=(), dtype=float32)\n",
            "Step:  70 train/avg_approx_loss_invalid : tf.Tensor(1.8934901, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  70 val/y_values_max : tf.Tensor(-0.80804974, shape=(), dtype=float32)\n",
            "Step:  70 val/y_values_mean : tf.Tensor(-48.484962, shape=(), dtype=float32)\n",
            "Step:  70 val/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  70 val/negatives_dist : tf.Tensor(-53.476532, shape=(), dtype=float32)\n",
            "Step:  70 val/mse_loss : tf.Tensor(930.3341, shape=(), dtype=float32)\n",
            "Step:  70 val/avg_approx_loss : tf.Tensor(2.1871474, shape=(), dtype=float32)\n",
            "Step:  70 val/avg_ranking_loss : tf.Tensor(-0.29353473, shape=(), dtype=float32)\n",
            "Step:  70 val/avg_ranking_train_loss : tf.Tensor(21.765175, shape=(), dtype=float32)\n",
            "Step:  70 val/avg_kendall_loss : tf.Tensor(-0.19185048, shape=(), dtype=float32)\n",
            "Step:  70 val/cql_loss : tf.Tensor(-53.476532, shape=(), dtype=float32)\n",
            "Step:  70 val/negatives_pred : tf.Tensor(-53.476532, shape=(), dtype=float32)\n",
            "Step:  70 val/model_pred_average : tf.Tensor(-55.997845, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  80 train/y_values_max : tf.Tensor(-0.9038583, shape=(), dtype=float32)\n",
            "Step:  80 train/y_values_mean : tf.Tensor(-47.936115, shape=(), dtype=float32)\n",
            "Step:  80 train/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  80 train/negatives_dist : tf.Tensor(-50.614216, shape=(), dtype=float32)\n",
            "Step:  80 train/mse_loss : tf.Tensor(785.9317, shape=(), dtype=float32)\n",
            "Step:  80 train/avg_approx_loss : tf.Tensor(2.087608, shape=(), dtype=float32)\n",
            "Step:  80 train/avg_ranking_loss : tf.Tensor(-0.057850055, shape=(), dtype=float32)\n",
            "Step:  80 train/avg_ranking_train_loss : tf.Tensor(16.756374, shape=(), dtype=float32)\n",
            "Step:  80 train/avg_kendall_loss : tf.Tensor(-0.032169104, shape=(), dtype=float32)\n",
            "Step:  80 train/cql_loss : tf.Tensor(-50.614216, shape=(), dtype=float32)\n",
            "Step:  80 train/negatives_pred : tf.Tensor(-50.614216, shape=(), dtype=float32)\n",
            "Step:  80 train/model_pred_average : tf.Tensor(-50.539318, shape=(), dtype=float32)\n",
            "Step:  80 train/invalid/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  80 train/y_value_infeasible : tf.Tensor(-50.425545, shape=(), dtype=float32)\n",
            "Step:  80 train/mse_loss_invalid : tf.Tensor(987.44763, shape=(), dtype=float32)\n",
            "Step:  80 train/mse_loss_overall : tf.Tensor(1773.3794, shape=(), dtype=float32)\n",
            "Step:  80 train/avg_approx_loss_invalid : tf.Tensor(2.67552, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  80 val/y_values_max : tf.Tensor(-0.80804974, shape=(), dtype=float32)\n",
            "Step:  80 val/y_values_mean : tf.Tensor(-50.924248, shape=(), dtype=float32)\n",
            "Step:  80 val/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  80 val/negatives_dist : tf.Tensor(-50.346638, shape=(), dtype=float32)\n",
            "Step:  80 val/mse_loss : tf.Tensor(828.5458, shape=(), dtype=float32)\n",
            "Step:  80 val/avg_approx_loss : tf.Tensor(2.064958, shape=(), dtype=float32)\n",
            "Step:  80 val/avg_ranking_loss : tf.Tensor(-0.11538826, shape=(), dtype=float32)\n",
            "Step:  80 val/avg_ranking_train_loss : tf.Tensor(18.36464, shape=(), dtype=float32)\n",
            "Step:  80 val/avg_kendall_loss : tf.Tensor(-0.064828455, shape=(), dtype=float32)\n",
            "Step:  80 val/cql_loss : tf.Tensor(-50.346638, shape=(), dtype=float32)\n",
            "Step:  80 val/negatives_pred : tf.Tensor(-50.346638, shape=(), dtype=float32)\n",
            "Step:  80 val/model_pred_average : tf.Tensor(-52.72599, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  90 train/y_values_max : tf.Tensor(-0.9038583, shape=(), dtype=float32)\n",
            "Step:  90 train/y_values_mean : tf.Tensor(-51.69957, shape=(), dtype=float32)\n",
            "Step:  90 train/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  90 train/negatives_dist : tf.Tensor(-48.031395, shape=(), dtype=float32)\n",
            "Step:  90 train/mse_loss : tf.Tensor(794.3457, shape=(), dtype=float32)\n",
            "Step:  90 train/avg_approx_loss : tf.Tensor(1.991349, shape=(), dtype=float32)\n",
            "Step:  90 train/avg_ranking_loss : tf.Tensor(-0.05236754, shape=(), dtype=float32)\n",
            "Step:  90 train/avg_ranking_train_loss : tf.Tensor(16.758024, shape=(), dtype=float32)\n",
            "Step:  90 train/avg_kendall_loss : tf.Tensor(-0.03039217, shape=(), dtype=float32)\n",
            "Step:  90 train/cql_loss : tf.Tensor(-48.031395, shape=(), dtype=float32)\n",
            "Step:  90 train/negatives_pred : tf.Tensor(-48.031395, shape=(), dtype=float32)\n",
            "Step:  90 train/model_pred_average : tf.Tensor(-48.295383, shape=(), dtype=float32)\n",
            "Step:  90 train/invalid/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  90 train/y_value_infeasible : tf.Tensor(-48.175266, shape=(), dtype=float32)\n",
            "Step:  90 train/mse_loss_invalid : tf.Tensor(927.43634, shape=(), dtype=float32)\n",
            "Step:  90 train/mse_loss_overall : tf.Tensor(1721.782, shape=(), dtype=float32)\n",
            "Step:  90 train/avg_approx_loss_invalid : tf.Tensor(1.9503968, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "-------------------------------------------------------\n",
            "Step:  90 val/y_values_max : tf.Tensor(-0.80804974, shape=(), dtype=float32)\n",
            "Step:  90 val/y_values_mean : tf.Tensor(-49.71262, shape=(), dtype=float32)\n",
            "Step:  90 val/vote_entropy : tf.Tensor(0.0, shape=(), dtype=float32)\n",
            "Step:  90 val/negatives_dist : tf.Tensor(-48.144753, shape=(), dtype=float32)\n",
            "Step:  90 val/mse_loss : tf.Tensor(898.52563, shape=(), dtype=float32)\n",
            "Step:  90 val/avg_approx_loss : tf.Tensor(1.6679305, shape=(), dtype=float32)\n",
            "Step:  90 val/avg_ranking_loss : tf.Tensor(-0.162691, shape=(), dtype=float32)\n",
            "Step:  90 val/avg_ranking_train_loss : tf.Tensor(19.78121, shape=(), dtype=float32)\n",
            "Step:  90 val/avg_kendall_loss : tf.Tensor(-0.102757335, shape=(), dtype=float32)\n",
            "Step:  90 val/cql_loss : tf.Tensor(-48.144753, shape=(), dtype=float32)\n",
            "Step:  90 val/negatives_pred : tf.Tensor(-48.144753, shape=(), dtype=float32)\n",
            "Step:  90 val/model_pred_average : tf.Tensor(-50.337814, shape=(), dtype=float32)\n",
            "-------------------------------------------------------\n",
            "Finished Training\n"
          ]
        }
      ],
      "source": [
        "#@title Running training\n",
        "\n",
        "# A toy example of running training PRIME. \n",
        "train_eval_offline(\n",
        "    config=config_str,\n",
        "    training_dataset=training_data,\n",
        "    validation_dataset=validation_data,\n",
        "    train_steps=100,\n",
        "    summary_freq=10,\n",
        "    eval_freq=10,\n",
        "    add_summary=True,\n",
        "    save_dir=None,\n",
        "    loss_type='mse+rank',\n",
        "    layers=(256, 256, 256),\n",
        "    with_ranking_penalty=True,\n",
        "    ranking_penalty_weight=0.01,\n",
        "    use_dropout=True,\n",
        "    cql_alpha=0.1,\n",
        "    infeasible_alpha=0.05\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m5KUS8rq3SQf"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "Z6W-wFO7FK0O",
        "ll-Ud813QwND"
      ],
      "name": "[PRIME, ICLR 2022] Colab Demonstration",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
